package com.pai4j.aigc.llm;

import com.pai4j.aigc.context.ContextKit;
import com.pai4j.aigc.context.service.ChatContextService;
import com.pai4j.aigc.llm.LlmSummarizerService;
import com.pai4j.aigc.context.model.ChatMessageContext;
import com.pai4j.aigc.llm.entity.LLMModelEntity;
import com.pai4j.aigc.llm.service.LLMModelService;
import com.pai4j.aigc.llm.service.LLMUsageService;
import com.pai4j.aigc.llm.service.MetricsService;
import com.pai4j.aigc.llm.service.PricingService;
import com.pai4j.aigc.llm.service.TokenEstimator;
import com.pai4j.common.enums.PromptTemplateEnum;
import com.pai4j.common.service.FreeMarkEngineService;
import com.pai4j.common.util.SpringContextUtil;
import com.pai4j.domain.vo.llm.ChatCompletionMessage;
import com.pai4j.domain.vo.llm.ChatCompletionResponse;
import com.pai4j.domain.vo.llm.ChatMessageRole;
import com.pai4j.domain.vo.llm.Usage;
import com.pai4j.domain.vo.request.request.mess.ChatMessageRequestVO;
import com.pai4j.domain.vo.response.UserBaseResponseInfoVO;
import com.pai4j.remote.user.UserServiceClient;
import com.fasterxml.jackson.databind.JsonNode;
import com.fasterxml.jackson.databind.ObjectMapper;
import lombok.extern.slf4j.Slf4j;
import org.apache.commons.lang3.StringUtils;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.beans.factory.annotation.Value;
import org.springframework.cloud.context.config.annotation.RefreshScope;
import org.springframework.stereotype.Service;
import org.springframework.web.servlet.mvc.method.annotation.SseEmitter;

import java.util.ArrayList;
import java.util.Arrays;
import java.util.HashMap;
import java.util.List;
import java.util.Map;
import java.util.Optional;
import java.util.UUID;

/**
 * gpt对话服务
 *
 * @author: CYM-pai
 * @date: 2025/05/11 09:20
 **/
@Slf4j
@RefreshScope
@Service
public class GPTChatService {

    @Autowired
    private UserServiceClient userServiceClient;

    @Autowired
    private FreeMarkEngineService freeMarkEngineService;

    @Value("${llm.config.model}")
    private String model;

    @Autowired
    private LLMUsageService usageService;

    @Autowired
    private TokenEstimator tokenEstimator;

    @Autowired
    private PricingService pricingService;

    @Autowired
    private MetricsService metricsService;

    @Autowired
    private LLMModelService modelService;

    /**
     * 没有设定角色时，使用默认角色
     */
    private static final String DEFAULT_GPT_ROLE = "你是一个智能助手AgentPAI，帮我解决一切问题。";
    private static final String DEFAULT_ROLE = "暂无签名";

    public String chat(SseEmitter sseEmitter, ChatMessageRequestVO chatMessage) {
        // 使用工厂模式，取对应的LLM服务，来完成对话
        AbstractLLMChatService abstractLLMChatService = LLMServiceFactory.getLLMService(model);
        // 获取用户发送消息内容
        String message = chatMessage.getMessage();
        // 取消息发送人、接收人信息
        String sendUserId = chatMessage.getSenderId();
        String receiverId = chatMessage.getReceiverId();
        UserBaseResponseInfoVO receiver = userServiceClient.getUserBaseInfo(receiverId).getData();
        String gptRoleSetting = receiver.getSignature();
        if (StringUtils.isBlank(gptRoleSetting) || DEFAULT_ROLE.equals(gptRoleSetting)) {
            gptRoleSetting = DEFAULT_GPT_ROLE;
        }
        // 确认/创建会话，并加载最近上下文（优先使用请求中的 sessionId，避免复用上一活跃会话）
        String sessionId = ContextKit.ensureSession(sendUserId, receiverId, chatMessage.getSessionId());
        List<ChatMessageContext> recent = ContextKit.loadRecent(sessionId, 20);
        Map<String, Object> vars = new HashMap<>();
        vars.put("roleSetting", gptRoleSetting);
        vars.put("content", message);
        String prompt;
        try {
            prompt = freeMarkEngineService.getContentByTemplate(PromptTemplateEnum.ROBOT_CHAT.getPath(), vars);
        } catch (Exception e) {
            log.error("GPTChatService.chat create prompt error. sendUserId:{}, receiverId:{}, msg:{}",
                    sendUserId, receiverId, message, e);
            return null;
        }
        final List<ChatCompletionMessage> messageList = new ArrayList<>();
        // 系统提示
        messageList.add(new ChatCompletionMessage(ChatMessageRole.SYSTEM.value(), gptRoleSetting));
        // 历史上下文
        if (recent != null && !recent.isEmpty()) {
            for (ChatMessageContext ctx : recent) {
                String role = ChatMessageRole.USER.value();
                if ("assistant".equalsIgnoreCase(ctx.getRole())) {
                    role = ChatMessageRole.ASSISTANT.value();
                } else if ("system".equalsIgnoreCase(ctx.getRole())) {
                    role = ChatMessageRole.SYSTEM.value();
                }
                messageList.add(new ChatCompletionMessage(role, ctx.getContent()));
            }
        }
        // 当前用户输入（模板后的prompt）
        messageList.add(new ChatCompletionMessage(ChatMessageRole.USER.value(), prompt));
        // Resolve model details for usage (fallback to global code only)
        Optional<LLMModelEntity> modelEntityOpt = modelService.findByGlobalCode(model);
        LLMModelEntity modelEntity = modelEntityOpt.orElse(null);
        String providerKey = modelEntity != null ? modelEntity.getProvider() : model;
        Long modelId = modelEntity != null ? modelEntity.getId() : -1L;
        String tenantId = (modelEntity != null) ? modelEntity.getTenantId() : null;
        long start = System.currentTimeMillis();
        String requestId = UUID.randomUUID().toString();
        try {
            // 仅将原始用户输入写入上下文（避免将包含“# 角色设定 ...”的模板prompt写入缓存）
            ContextKit.appendUser(sessionId, message);
            // Build runtime config from model entity if available (apiKey/model/baseUrl + parameters)
            String result;
            if (modelEntity != null) {
                Integer maxTokens = null;
                Float temperature = null;
                try {
                    String params = modelEntity.getParameters();
                    if (params != null && !params.isEmpty()) {
                        ObjectMapper mapper = new ObjectMapper();
                        JsonNode root = mapper.readTree(params);
                        if (root.has("maxTokens") && !root.get("maxTokens").isNull()) {
                            maxTokens = root.get("maxTokens").asInt();
                        }
                        if (root.has("temperature") && !root.get("temperature").isNull()) {
                            temperature = (float) root.get("temperature").asDouble();
                        }
                    }
                } catch (Exception ignore) {}
                LLMRuntimeConfig runtime = new LLMRuntimeConfig(
                        modelEntity.getApiKey(),
                        modelEntity.getModel(),
                        modelEntity.getBaseUrl(),
                        maxTokens,
                        temperature
                );
                result = abstractLLMChatService.chat(sseEmitter, messageList, runtime);
            } else {
                result = abstractLLMChatService.chat(sseEmitter, messageList);
            }
            if (StringUtils.isNotBlank(result)) {
                // 写入助手回复
                ContextKit.appendAssistant(sessionId, result);
                // 首轮后自动生成标题（若为空），LLM生成
                LlmSummarizerService titleSvc = SpringContextUtil.getBean(LlmSummarizerService.class);
                String title = titleSvc.generateTitle(message, result);
                // 顺便触发一次阈值检查
                ContextKit.summarizeIfNeeded(sessionId);
                ChatContextService svc = SpringContextUtil.getBean(ChatContextService.class);
                svc.setTitleIfEmpty(sessionId, title);
            }
            // Record usage for streaming: prefer provider-reported usage; fallback to estimation
            long latency = System.currentTimeMillis() - start;
            Integer pt;
            Integer ct;
            Integer total;
            Usage providerUsage = abstractLLMChatService.consumeLastUsage();
            if (providerUsage != null && providerUsage.getTotalTokens() > 0) {
                pt = providerUsage.getPromptTokens();
                ct = providerUsage.getCompletionTokens();
                total = providerUsage.getTotalTokens();
            } else {
                pt = tokenEstimator.estimateTokens(message);
                ct = tokenEstimator.estimateTokens(result);
                total = pt + ct;
            }
            Integer costCents = pricingService.calcCostCents(modelEntity, pt, ct);
            usageService.recordSuccess(
                    tenantId,
                    requestId,
                    sendUserId,
                    modelId,
                    providerKey,
                    modelEntity != null ? modelEntity.getCode() : providerKey,
                    pt,
                    ct,
                    total,
                    latency,
                    costCents
            );
            metricsService.recordSuccess(providerKey, tenantId, modelEntity != null ? modelEntity.getCode() : providerKey, latency);
            return result;
        } catch (Exception e) {
            long latency = System.currentTimeMillis() - start;
            usageService.recordFailure(
                    tenantId,
                    requestId,
                    sendUserId,
                    modelId,
                    providerKey,
                    modelEntity != null ? modelEntity.getCode() : providerKey,
                    "LLM_ERROR",
                    e.getMessage(),
                    latency
            );
            metricsService.recordFailure(providerKey, tenantId, modelEntity != null ? modelEntity.getCode() : providerKey, latency);
            log.error("GPTChatService.chat error. sendUserId:{}, receiverId:{}, msg:{}",
                    sendUserId, receiverId, message, e);
        }
        return null;
    }
}
